#根据预测出来的txt文件裁剪图片
import os
import cv2
import math
from tqdm import tqdm

image_input = '/home/xys/CloundShiProjects/traffic_light/trafficlight_classify/qixian_test/jpg/'
txt_input = '/home/xys/CloundShiProjects/traffic_light/trafficlight_classify/qixian_test/txt/'
path_output = "/home/xys/CloundShiProjects/traffic_light/trafficlight_classify/qixian_test/traffic_light/"  #裁剪出来的小图保存的根目录

img_total = []
txt_total = []

file_image = os.listdir(image_input)
for filename in file_image:
     first,last = os.path.splitext(filename)
     img_total.append(first)

file_txt = os.listdir(txt_input)
for filename in file_txt:
     first,last = os.path.splitext(filename)
     txt_total.append(first)

for img_ in tqdm(img_total):
    if img_ in txt_total:
        filename_img = img_+".jpg"
        path1 = os.path.join(image_input,filename_img)
        img = cv2.imread(path1)
        filename_txt = img_     #预测出来的txt文件没有后缀名，有则加 {+".txt"}
        try:
            h = img.shape[0]
        except:
            import shutil
            shutil.move(path1, '/home/xys/CloundShiProjects/traffic_light/trafficlight_classify/qixian_test/')
            continue
        # h = img.shape[0]
        w = img.shape[1]
        n = 1
        with open(os.path.join(txt_input,filename_txt),"r+",encoding="utf-8",errors="ignore") as f:
            for line in f:
                aa = line.split(" ")
                if float(aa[5]) < 0.5: continue
                # if not (int(aa[0]) == 1 or int(aa[0]) == 0): continue
                x_center = w * float(aa[1])       # aa[1]左上点的x坐标
                y_center = h * float(aa[2])       # aa[2]左上点的y坐标
                # width = int(w * float(aa[3]))  # aa[3]图片width
                # width = int(float(width) * 1.3)
                # height = int(h * float(aa[4]))  # aa[4]图片height
                # height = int(float(height) * 1.3)
                # lefttopx = int(x_center - width / 2.0)
                # lefttopy = int(y_center - height / 2.0)
                width = int(w*float(aa[3]))       # aa[3]图片width
                # width = int(float(width) * 1)
                height = int(h*float(aa[4]))      # aa[4]图片height
                # height = int(float(height) * 1)
                lefttopx = int(math.fabs(x_center-width/2.0))
                lefttopy = int(math.fabs(y_center-height/2.0))
                # roi = img[lefttopy+1:lefttopy+height+3,lefttopx+1:lefttopx+width+1]   # [左上y:右下y,左上x:右下x] (y1:y2,x1:x2)需要调参，否则裁剪出来的小图可能不太好
                roi = img[lefttopy:lefttopy + height, lefttopx:lefttopx + width]
                if roi.size == 0: continue
                filename_last = img_+"_0"+str(n)+".jpg"      # 裁剪出来的小图文件名
                path2 = path_output          # 需要在path_output路径下创建一个cut_txt文件夹
                # print('path2:', path2)                    # 裁剪小图的保存位置
                if not os.path.exists(os.path.join(path2,str(aa[0]))):
                    os.mkdir(os.path.join(path2,str(aa[0])))
                cv2.imwrite(os.path.join(path2,str(aa[0]) ,filename_last),roi)
                n = n+1
    else:
        continue
